# text2sql_query.py
import os, re
import logging
import yaml
from openai import OpenAI
from dotenv import load_dotenv
from pymilvus import connections, utility, Collection
from sentence_transformers import SentenceTransformer
from sqlalchemy import create_engine, text

# 1. 环境与日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
load_dotenv()  # 加载 .env 环境变量

# 2. 初始化 OpenAI API（使用最新 Response API）
#openai.api_key = os.getenv("OPENAI_API_KEY", "xxxxx")

# 建议使用新 Response API 风格
# 例如: openai.chat.completions.create(...) 而非旧的 ChatCompletion.create

#MODEL_NAME = os.getenv("OPENAI_MODEL", "deepseek-chat")

# 3. 嵌入函数初始化 - 使用本地BGE-M3模型
bge_m3_model = "D:\\code\\aicode\\sentence-transformers\\bge-m3"
embedding_model = SentenceTransformer(bge_m3_model)

def local_embedding_function(texts):
    """本地嵌入函数包装器"""
    if isinstance(texts, str):
        texts = [texts]
    embeddings = embedding_model.encode(texts, 
                                      batch_size=12, 
                                      convert_to_tensor=False,
                                      normalize_embeddings=True)
    return embeddings.tolist()

# 4. Milvus 客户端连接 - 使用Docker部署的Milvus服务
connections.connect(
    alias="default",
    host='localhost',     # Milvus服务器地址
    port='19530'          # Milvus服务器端口
)
# 检查连接状态
logging.info(f"[连接] 成功连接到Milvus服务器: {utility.list_collections()}")

# 5. 嵌入函数实例化
# embedding_fn = init_embedding()

# 6. 数据库连接（SAKILA）
DB_URL = os.getenv(
    "SAKILA_DB_URL", 
    "mysql+pymysql://root:123456@localhost:3306/sakila"
)
engine = create_engine(DB_URL)



def extract_entity_from_string(hit_string):
    """
    Extract entity dictionary from string representation
    
    Args:
        hit_string: String representation of a search result
        
    Returns:
        Dictionary containing entity content
    """
    # Extract entity content using regex
    entity_pattern = r"entity: ({.*})"
    entity_match = re.search(entity_pattern, hit_string, re.DOTALL)
    
    if entity_match:
        entity_str = entity_match.group(1)
        return eval(entity_str)
    else:
        return {}


# 7. 检索函数
def retrieve(collection_name: str, query_emb: list, top_k: int = 3, fields: list = None):
    collection = Collection(collection_name)
    collection.load()
    
    search_params = {
        "metric_type": "COSINE",
        "params": {"nprobe": 10}
    }
    
    results = collection.search(
        data=[query_emb],
        anns_field="vector",
        param=search_params,
        limit=top_k,
        output_fields=fields or []
    )
    logging.info(f"[检索] {collection_name} 检索结果数量: {len(results[0])}")

    hits = []
    for entity in results[0]:
        hit = extract_entity_from_string(str(entity))
        hits.append(hit)

    return hits



# 8. SQL 提取函数
def extract_sql(text: str) -> str:
    # 尝试匹配 SQL 代码块
    sql_blocks = re.findall(r'```sql\n(.*?)\n```', text, re.DOTALL)
    if sql_blocks:
        return sql_blocks[0].strip()
    
    # 如果没有找到代码块，尝试匹配 SELECT 语句
    select_match = re.search(r'SELECT.*?;', text, re.DOTALL)
    if select_match:
        return select_match.group(0).strip()
    
    # 如果都没有找到，返回原始文本
    return text.strip()



# 9. 核心流程：自然语言 -> SQL -> 执行 -> 返回
def text2sql(question: str):
    # 8.1 用户提问嵌入
    q_emb = local_embedding_function([question])[0]
    logging.info(f"[检索] 问题嵌入完成")

    # 8.2 RAG 检索：DDL
    ddl_hits = retrieve("ddl_knowledge", q_emb, top_k=3, fields=["ddl_text"])
    logging.info(f"[检索] DDL检索结果: {ddl_hits}")

    try:
        ddl_context = "\n".join(hit.get("ddl_text", "") for hit in ddl_hits)
    except Exception as e:
        logging.error(f"[检索] DDL处理错误: {e}")
        ddl_context = ""
    logging.info(f"[检索] DDL结果: {ddl_context}")


    # 8.3 RAG 检索：示例对
    q2sql_hits = retrieve("q2sql_knowledge", q_emb, top_k=3, fields=["question", "sql_text"])
    logging.info(f"[检索] Q2SQL检索结果: {q2sql_hits}")
    try:
        example_context = "\n".join(
            f"NL: \"{hit.get('question', '')}\"\nSQL: \"{hit.get('sql_text', '')}\"" 
            for hit in q2sql_hits
        )
    except Exception as e:
        logging.error(f"[检索] Q2SQL处理错误: {e}")
        example_context = ""

    # 8.4 RAG 检索：字段描述
    desc_hits = retrieve("dbdesc_knowledge", q_emb, top_k=5, fields=["table_name", "column_name", "description"])
    logging.info(f"[检索] 字段描述检索结果: {desc_hits}")
    try:
        desc_context = "\n".join(
            f"{hit.get('table_name', '')}.{hit.get('column_name', '')}: {hit.get('description', '')}"
            for hit in desc_hits
        )
    except Exception as e:
        logging.error(f"[检索] 字段描述处理错误: {e}")
        desc_context = ""

    # 8.5 Prompt 组装
    prompt = (
        f"### Schema Definitions:\n{ddl_context}\n"
        f"### Field Descriptions:\n{desc_context}\n"
        f"### Examples:\n{example_context}\n"
        f"### Query:\n\"{question}\"\n"
        "请只返回SQL语句，不要包含任何解释或说明。"
    )
    logging.info("[生成] 开始生成SQL")

    # 8.6 调用最新 Response API
    openai_client = OpenAI(api_key="xxxxx", base_url="https://api.deepseek.com")
    response = openai_client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": prompt}],
        stream=False
    )
    raw_sql = response.choices[0].message.content.strip()
    sql = extract_sql(raw_sql)
    logging.info(f"[生成] 原始输出: {raw_sql}")
    logging.info(f"[生成] 提取的SQL: {sql}")

    # 8.7 执行并打印结果
    try:
        with engine.connect() as conn:
            result = conn.execute(text(sql))
            cols = result.keys()
            rows = result.fetchall()
            print("\n查询结果：")
            print("列名：", cols)
            for r in rows:
                print(r)
    except Exception as e:
        logging.error(f"[执行] 执行失败: {e}")
        print("执行错误：", e)

# 9. 程序入口
if __name__ == "__main__":
    user_q = input("请输入您的自然语言查询： ")
    text2sql(user_q)